package com.hy.ratelimiting.limiting;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.concurrent.atomic.AtomicLong;

/**
 * Created With IntelliJ IDEA.
 * Descriptions:
 * <p>
 * User:Mr.Du
 * Date:2024/4/29
 * Time:11:35
 */

@Slf4j
@Service
public class TokenBucketRateLimiter {

    private long capacity = 10; // 令牌桶容量
    private long rate = 10; // 令牌生成速率，单位：毫秒
    private AtomicLong tokens = new AtomicLong(10); // 当前令牌数量
    private long lastRefillTime; // 上次令牌补充时间

    public TokenBucketRateLimiter() {
    }

    public TokenBucketRateLimiter(long capacity, long rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = new AtomicLong(capacity); // 初始时，令牌桶满
        this.lastRefillTime = System.currentTimeMillis(); // 初始化时记录上次补充时间
    }

    // 尝试获取令牌，成功返回true，失败返回false
    public synchronized boolean allowRequest() {
        refillTokens(); // 补充令牌
        long currentTokens = tokens.get(); // 获取当前令牌数量
        if (currentTokens > 0) { // 令牌桶中有令牌
            tokens.decrementAndGet(); // 消费一个令牌
            return true; // 获取令牌成功
        }
        return false; // 令牌桶中无令牌，获取失败
    }

    // 补充令牌
    private synchronized void refillTokens() {
        long now = System.currentTimeMillis(); // 当前时间
        long elapsedTime = now - lastRefillTime; // 距离上次补充令牌的时间间隔
        log.info("当前时间：" + now + ", 距离上次补充令牌的时间间隔：" + elapsedTime);
        long tokensToAdd = elapsedTime * rate / 1000; // 计算需要补充的令牌数量
        log.info("当前令牌数量：" + tokens.get() + ", 补充令牌数量：" + tokensToAdd);
        if (tokensToAdd > 0) { // 需要补充令牌
            tokens.set(Math.min(capacity, tokens.get() + tokensToAdd)); // 补充令牌，并确保不超过容量
            lastRefillTime = now; // 更新上次补充时间
        }
    }

}
